import numpy as np
from sklearn.neighbors import NearestNeighbors

# 模拟向量数据
np.random.seed(42)
data_vectors = np.random.rand(100, 5)  # 100个5维向量

# 实现倒排索引
class InvertedIndex:
    def __init__(self):
        self.index = {}
    
    def add_to_index(self, vector_id, bucket):
        """将向量ID添加到倒排索引对应的桶"""
        if bucket not in self.index:
            self.index[bucket] = []
        self.index[bucket].append(vector_id)
    
    def search(self, bucket):
        """检索某个桶中的所有向量ID"""
        return self.index.get(bucket, [])

# 构建倒排索引
inverted_index = InvertedIndex()
for i, vector in enumerate(data_vectors):
    bucket = int(vector[0] * 10)  # 简单分桶: 根据第一个维度值划分
    inverted_index.add_to_index(i, bucket)

# 检索倒排索引
bucket_to_search = 2
result_ids = inverted_index.search(bucket_to_search)
print(f"倒排索引中桶 {bucket_to_search} 的向量ID:", result_ids)


# 实现图索引
class GraphIndex:
    def __init__(self, data_vectors):
        self.data_vectors = data_vectors
        self.graph = None
    
    def build_graph(self, k=5):
        """使用sklearn构建k近邻图"""
        nbrs = NearestNeighbors(n_neighbors=k, algorithm='auto').fit(self.data_vectors)
        distances, indices = nbrs.kneighbors(self.data_vectors)
        self.graph = {i: indices[i].tolist() for i in range(len(self.data_vectors))}
    
    def search(self, query_vector, top_k=5):
        """从图中找到与查询向量最相似的k个邻居"""
        distances = [np.linalg.norm(query_vector - vec) for vec in self.data_vectors]
        nearest_indices = np.argsort(distances)[:top_k]
        return nearest_indices

# 构建图索引
graph_index = GraphIndex(data_vectors)
graph_index.build_graph()

# 检索图索引
query_vector = np.array([0.5, 0.5, 0.5, 0.5, 0.5])  # 查询向量
nearest_neighbors = graph_index.search(query_vector, top_k=3)
print("图索引中最相似的向量ID:", nearest_neighbors)

class Partitioning:
    def __init__(self, data_vectors, num_partitions):
        self.data_vectors = data_vectors
        self.num_partitions = num_partitions
        self.partitions = {i: [] for i in range(num_partitions)}
    
    def assign_to_partitions(self):
        """根据第一个维度的值将向量分区"""
        for i, vector in enumerate(self.data_vectors):
            partition_id = int(vector[0] * self.num_partitions)
            self.partitions[partition_id].append(i)
    
    def search_in_partition(self, partition_id, query_vector, top_k=5):
        """在指定分区中检索最相似的向量"""
        if partition_id not in self.partitions or not self.partitions[partition_id]:
            return []
        partition_vectors = [self.data_vectors[i] for i in self.partitions[partition_id]]
        distances = [np.linalg.norm(query_vector - vec) for vec in partition_vectors]
        nearest_indices = np.argsort(distances)[:top_k]
        return [self.partitions[partition_id][i] for i in nearest_indices]

# 创建分区
num_partitions = 10
partitioning = Partitioning(data_vectors, num_partitions)
partitioning.assign_to_partitions()

# 在分区中检索
partition_to_search = 3
nearest_in_partition = partitioning.search_in_partition(
    partition_to_search, query_vector, top_k=3)
print(f"分区 {partition_to_search} 中最相似的向量ID:", nearest_in_partition)